package com.zst.microservicenodeone.hystrix;

import com.netflix.hystrix.HystrixThreadPoolKey;
import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategy;
import com.netflix.hystrix.strategy.concurrency.HystrixRequestVariable;
import com.netflix.hystrix.strategy.concurrency.HystrixRequestVariableLifecycle;
import com.netflix.hystrix.strategy.properties.HystrixProperty;
import com.zst.microservicenodeone.context.UserContext;
import com.zst.microservicenodeone.context.UserContextHolder;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

public class ThreadLocalStrategy extends HystrixConcurrencyStrategy {
    private HystrixConcurrencyStrategy existed;

    public ThreadLocalStrategy(HystrixConcurrencyStrategy existed){
        this.existed = existed;
    }

    @Override
    public BlockingQueue<Runnable> getBlockingQueue(int maxQueueSize){
        return existed == null ? super.getBlockingQueue(maxQueueSize) : existed.getBlockingQueue(maxQueueSize);
    }

    @Override
    public <T>HystrixRequestVariable<T> getRequestVariable(HystrixRequestVariableLifecycle<T> rv){
        return existed != null ? existed.getRequestVariable(rv) : super.getRequestVariable(rv);
    }

    @Override
    public ThreadPoolExecutor getThreadPool(
            HystrixThreadPoolKey threadPoolKey,
            HystrixProperty<Integer> corePoolSize,
            HystrixProperty<Integer> maximumPoolSize,
            HystrixProperty<Integer> keepAliveTime,
            TimeUnit unit,
            BlockingQueue<Runnable> workQueue
    ){
        return existed != null ? existed.getThreadPool(threadPoolKey, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue)
                : super.getThreadPool(threadPoolKey, corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    @Override
    public <T> Callable<T> wrapCallable(Callable<T> callable){
        return existed != null ? existed.wrapCallable(new WrappedCallable(callable, UserContextHolder.getContext()))
                : super.wrapCallable(new WrappedCallable(callable, UserContextHolder.getContext()));
    }

    static final class WrappedCallable<T> implements Callable<T> {
        private final Callable<T> target;
        private UserContext userContext;

        public WrappedCallable(Callable<T> target, UserContext userContext){
            this.target = target;
            this.userContext = userContext;
        }

        @Override
        public T call() throws Exception {
            try{
                UserContextHolder.setContext(userContext);
                return target.call();
            }finally {
                this.userContext = null;
            }
        }
    }
}
